Skip to content

Joint passage size #28

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Joint passage size #28

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
eca62e3
changes for joint passage size
May 24, 2024
05b0a4f
WIP agent pos shaping
kfu02 May 28, 2024
3b01502
Revert "WIP agent pos shaping"
kfu02 May 29, 2024
84eb807
Merge branch 'main' into joint-passage-size
kfu02 May 29, 2024
4028402
modify done(), modify controller params
kfu02 May 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 48 additions & 18 deletions vmas/scenarios/joint_passage_size.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):

assert self.n_passages == 3 or self.n_passages == 4

# NOTE: I changed how our VMAS viewer works s.t. it requires a world_semidim, my bad
self.world_semidim = 1.0

self.plot_grid = False

# Make world
Expand Down Expand Up @@ -100,7 +103,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
self.n_boxes = int(self.scenario_length // self.passage_length)
self.min_collision_distance = 0.005

cotnroller_params = [2.0, 10, 0.00001]
controller_params = [1.0, 0.0, 0.0]

# Add agents
agent = Agent(
Expand All @@ -112,7 +115,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
f_range=10,
)
agent.controller = VelocityController(
agent, world, cotnroller_params, "standard"
agent, world, controller_params, "standard"
)
world.add_agent(agent)
agent = Agent(
Expand All @@ -126,7 +129,7 @@ def make_world(self, batch_dim: int, device: torch.device, **kwargs):
f_range=10,
)
agent.controller = VelocityController(
agent, world, cotnroller_params, "standard"
agent, world, controller_params, "standard"
)
world.add_agent(agent)

Expand Down Expand Up @@ -510,6 +513,20 @@ def observation(self, agent: Agent):
angle_to_vector(self.goal.state.rot),
] + ([angle_to_vector(joint_angle)] if self.observe_joint_angle else [])

# NOTE: I add the size capability in "mixed" form manually here:
radius = agent.shape.radius
mean_radius = (self.agent_radius + self.agent_radius_2) / 2
relative_radius = agent.shape.radius - mean_radius
capability_repr = [
torch.tensor(
radius, device=self.world.device
).repeat(self.world.batch_dim, 1),
torch.tensor(
relative_radius, device=self.world.device
).repeat(self.world.batch_dim, 1),
]
observations += capability_repr

if self.obs_noise > 0:
for i, obs in enumerate(observations):
noise = torch.zeros(
Expand All @@ -526,25 +543,35 @@ def observation(self, agent: Agent):
)

def done(self):
return torch.all(
(
torch.linalg.vector_norm(
self.joint.landmark.state.pos - self.goal.state.pos, dim=1
)
<= 0.01
)
* (
get_line_angle_dist_0_180(
self.joint.landmark.state.rot, self.goal.state.rot
).unsqueeze(-1)
<= 0.01
),
dim=1,
)
# reimplementation
pos_done = torch.linalg.vector_norm(self.joint.landmark.state.pos - self.goal.state.pos, dim=1) <= 0.01
rot_done = get_line_angle_dist_0_180(self.joint.landmark.state.rot, self.goal.state.rot) <= 0.01
done = torch.logical_and(pos_done, rot_done)
return done

# original
# return torch.all(
# (
# torch.linalg.vector_norm(
# self.joint.landmark.state.pos - self.goal.state.pos, dim=1
# )
# <= 0.01
# )
# * (
# get_line_angle_dist_0_180(
# self.joint.landmark.state.rot, self.goal.state.rot
# ) # .unsqueeze(-1)
# <= 0.01
# ),
# dim=1,
# )

def info(self, agent: Agent) -> Dict[str, Tensor]:
is_first = self.world.agents[0] == agent
if is_first:
dist_to_goal = torch.linalg.vector_norm(self.joint.landmark.state.pos - self.goal.state.pos, dim=1)
rot_to_goal = get_line_angle_dist_0_180(self.joint.landmark.state.rot, self.goal.state.rot)

just_passed = self.all_passed * (self.passed == 0)
self.passed[just_passed] = 100
self.info_stored = {
Expand All @@ -553,6 +580,9 @@ def info(self, agent: Agent) -> Dict[str, Tensor]:
"collision_rew": self.collision_rew,
"energy_rew": self.energy_rew,
"passed": just_passed.to(torch.int),
"success_rate": self.done(), # wandb averages over n_envs
"dist_to_goal": dist_to_goal,
"rot_to_goal": rot_to_goal,
}
return self.info_stored

Expand Down